# -*- coding: utf-8 -*-

import os
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
from sklearn.metrics.pairwise import cosine_similarity
from tensorflow.keras.applications.vgg16 import VGG16, preprocess_input
from tensorflow.keras.preprocessing import image
from tensorflow.keras.models import Model

# 加载预训练的VGG16模型，不包含顶部的全连接层
base_model = VGG16(weights='imagenet', include_top=False)
# 创建特征提取模型，使用最后一个卷积层的输出作为特征
model = Model(inputs=base_model.input, outputs=base_model.get_layer('block5_pool').output)
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]

def extract_features(img_path):
    """提取图片特征"""
    try:
        # 加载图片并调整大小
        img = image.load_img(img_path, target_size=(224, 224))
        # 转换为数组
        img_array = image.img_to_array(img)
        # 添加批次维度
        img_array = np.expand_dims(img_array, axis=0)
        # 预处理图片
        img_preprocessed = preprocess_input(img_array)
        # 提取特征
        features = model.predict(img_preprocessed)
        # 展平特征数组
        return features.flatten()
    except Exception as e:
        print(f"处理图片 {img_path} 时出错: {e}")
        return None

def find_similar_images(query_img_path, database_dir, top_n=5):
    """查找相似图片"""
    # 提取查询图片的特征
    query_features = extract_features(query_img_path)
    if query_features is None:
        return []
    
    # 提取数据库中所有图片的特征
    database_features = {}
    for filename in os.listdir(database_dir):
        if filename.lower().endswith(('.png', '.jpg', '.jpeg','.webp')):
            img_path = os.path.join(database_dir, filename)
            features = extract_features(img_path)
            if features is not None:
                database_features[img_path] = features
    
    # 计算相似度
    similarities = {}
    for img_path, features in database_features.items():
        # 计算余弦相似度
        similarity = cosine_similarity([query_features], [features])[0][0]
        similarities[img_path] = similarity
    
    # 按相似度排序并返回前top_n张图片
    similar_images = sorted(similarities.items(), key=lambda x: x[1], reverse=True)[:top_n]
    return similar_images

def display_similar_images(query_img_path, similar_images):
    """显示查询图片和相似图片"""
    plt.figure(figsize=(15, 10))
    
    # 显示查询图片
    plt.subplot(1, len(similar_images) + 1, 1)
    plt.title("查询图片")
    plt.imshow(Image.open(query_img_path))
    plt.axis('off')
    
    # 显示相似图片
    for i, (img_path, similarity) in enumerate(similar_images, 2):
        plt.subplot(1, len(similar_images) + 1, i)
        plt.title(f"相似度: {similarity:.4f}")
        plt.imshow(Image.open(img_path))
        plt.axis('off')
    
    plt.tight_layout()
    plt.show()

# 使用示例
if __name__ == "__main__":
    # 设置查询图片路径和图片数据库目录
    query_image_path = "D:\\Images\\fish2.png"
    image_database_directory = "D:\\Images"
    
    # 查找相似图片
    similar_images = find_similar_images(query_image_path, image_database_directory)
    
    # 显示结果
    if similar_images:
        display_similar_images(query_image_path, similar_images)
    else:
        print("未找到相似图片或处理过程中出错。")    